
//% color="#0062AC" iconWidth=50 iconHeight=40
namespace librosa_spec {

    //% block="录音[t]秒直到结束，保存到文件[r_path]" blockType="command"
    //% t.shadow="number"  t.defl=2
    //% r_path.shadow="string" r_path.defl="record.wav"
    export function wav_record(parameter: any, block: any) {
        let t = parameter.t.code;
        let r_path = parameter.r_path.code;
    Generator.addImport(`\nimport sounddevice as sd\nimport soundfile as sf \n`)
    Generator.addCode(`
# 设置参数
samplerate = 16000  # 采样率
duration = ${t}  # 捕捉声音的时间长度（秒）
output_file = ${r_path}

# 捕捉声音
print("开始捕捉声音，请说话...")

audio_data = sd.rec(int(duration * samplerate), samplerate=samplerate, channels=2, dtype='float32')
sd.wait()  # 等待录制结束
print("声音捕捉完成。")

# 保存为 WAV 文件
sf.write(output_file, audio_data, samplerate)
print(f"音频已保存为 {output_file}")
`)
}

    //% block="音转图加载准备" blockType="command"
    export function colorhelp(parameter: any, block: any) {
        Generator.addImport(`\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib.cm import get_cmap`)
        Generator.addCode(`
viridis_cmap = plt.get_cmap('viridis')                   
color_map = viridis_cmap.colors                          
color_map = (np.array(color_map) * 255).astype(np.uint8) `)

        }
        //% block="---"
        export function noteSep1() {

        }
        //% block="单音频处理" blockType="tag"
        export function tagtest1() {}


    //% block="导入音频 路径[path] 赋值给序列[wave]" blockType="command"
    //% path.shadow="string" path.defl="sound.wav"
    //% wave.shadow="normal" wave.defl="wave"
    export function InputSound(parameter: any, block: any) {
        let wave=parameter.wave.code;
        let path=parameter.path.code;
        Generator.addImport(`\nimport librosa\n`)
        Generator.addCode(`path = ${path}
${wave},fs = librosa.load(path, sr = 16000)`)

    }
    //% block="导出序列[wave] 响度振幅图" blockType="command"
    //% wave.shadow="normal" wave.defl="wave"
    export function loudout(parameter: any, block: any) {
        let wave=parameter.wave.code;

        Generator.addImport(`\nfrom torchvision import transforms\n`)
        Generator.addCode(`
# 计算响度（以分贝为单位）
loudness = librosa.feature.rms(y=${wave})**2  # RMS的平方，因为响度通常与功率成正比

# 归一化响度数据到0-255的范围
loudness_normalized = (loudness - np.min(loudness)) / (np.max(loudness) - np.min(loudness)) * 255
loudness_normalized = loudness_normalized.astype(np.uint8)  # 转换为整数

h,w = loudness_normalized.shape
rgb_matrix = np.array([color_map[i] for i in loudness_normalized.flatten()]).reshape(h, w, 3)
loud_rgb = rgb_matrix/255
test = transforms.ToTensor()(loud_rgb)
test = transforms.Resize((224, 224))(test)
plt.imsave(path[:-4]+'_loudness.jpg', np.array(test).transpose(1,2,0))
print('数据保存成功')`)

    }


    //% block="导出序列[wave] 梅尔频谱图" blockType="command"
    //% wave.shadow="normal" wave.defl="wave"
    export function sigout(parameter: any, block: any) {
        let wave=parameter.wave.code;

        Generator.addImport(`\nfrom torchvision import transforms\n`)
        Generator.addCode(`
spc = librosa.feature.melspectrogram(y=${wave}, sr=fs, n_fft=512)
print(spc)
#转化为log形式
spc = librosa.power_to_db(spc, ref=np.max)
#数据归一化
spec_new = (((spc+80)/80)*255).astype(np.uint8)
h,w = spec_new.shape
rgb_matrix = np.array([color_map[i] for i in spec_new.flatten()]).reshape(h, w, 3)
spec_rgb = rgb_matrix/255
test = transforms.ToTensor()(spec_rgb)
test = transforms.Resize((224, 224))(test)
plt.imsave(path[:-4]+'_melspc.jpg', np.array(test).transpose(1,2,0))
print('数据保存成功')`)

    }
        //% block="---"
        export function noteSep2() {

        }
        //% block="批量处理" blockType="tag"
        export function tagtest2() {}

    //% block="导入音频素材文件夹 路径[PATH] 每段时长[time]秒" blockType="command"
    //% PATH.shadow="string" PATH.defl="root/原音频素材/打开"
    //% time.shadow="number" time.defl="2"
    export function Input_Dir(parameter: any, block: any) {
        let PATH=parameter.PATH.code;
        let time = parameter.time.code;

        Generator.addImport(`import os\nimport librosa\nimport glob\n`)
        Generator.addCode(`s_path = ${PATH}\ntime=${time}
X =[]
y= []
# Load the file
audio_files = []

for root, dirs, files in os.walk(s_path):
    for file in files:
        if file.endswith(('.mp3', '.wav', '.flac', '.m4a')):
            audio_files.append(os.path.join(root, file))
# read file 
for file in audio_files:  
    # 获得标签
    y.append(file.split('\\\\')[-1][:-4])
    # 读取音频
    wave,fs = librosa.load(file, sr = 16000)
    # time对应的样点数
    sample = int(time * fs)
    # 若音频短于time的样点数，则在音频后面补零，若音频长于time的样点数，则对音频截断
    if wave.size <= sample:
        wave = np.concatenate((wave,np.array((sample - wave.size) * [0])))
    else:
        wave = wave[0:sample]
    # 获取梅尔频谱
    spc = librosa.feature.melspectrogram(y=wave, sr=fs, n_fft=512)
    #转化为log形式
    spc = librosa.power_to_db(spc, ref=np.max)
    #数据归一化
    spec_new = (((spc+80)/80)*255).astype(np.uint8)
    h,w = spec_new.shape
    rgb_matrix = np.array([color_map[i] for i in spec_new.flatten()]).reshape(h, w, 3)
    spec_rgb = rgb_matrix/255
    # 添加到数据
    X.append(spec_rgb)  
print(y)
`)

    }

    //% block="导出一组梅尔频谱到路径[output]" blockType="command"
    //% output.shadow="string" output.defl="root/频谱图/打开"
    export function imgoutput(parameter: any, block: any) {
        let output=parameter.output.code;

        Generator.addImport(`\nfrom torchvision import transforms\n`)
        Generator.addCode(`out_path = ${output}
for i,j in enumerate(y):
    if out_path[-1]!='/':
        out_path+='/'
    test = X[i]
    test = transforms.ToTensor()(test)
    test = transforms.Resize((224, 224))(test)
    plt.imsave(out_path+str(i)+'.jpg', np.array(test).transpose(1,2,0))
print('数据保存成功')
            `)

}
}